from langchain_core.runnables import RunnableWithMessageHistory,ConfigurableFieldSpec
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.chat_message_histories import ChatMessageHistory

import os
# 定义模型
llm = ChatOpenAI(model_name="deepseek-chat", api_key=os.environ["DEEPSEEK_API_KEY"],base_url=os.environ["DEEPSEEK_BASE_URL"])

# 定义提示词模板
prompt_template = ChatPromptTemplate.from_messages([
    ('system','你是一个乐于助人的助手，请尽你所能回答问题'),
    MessagesPlaceholder(variable_name="history_msg"), # 获取用户的聊天记录  variable_name 为 RunnableWithMessageHistory中的input_messages_key
    ('human',"{input_msg}")
])

# 创建一个runable
chain = prompt_template | llm

#保存聊天记录
chat_history = {}

"""
获取聊天记录
"""
def get_chat_history(user_id:str,conversation_id:str):
    if (user_id,conversation_id) not in chat_history:
        chat_history[(user_id,conversation_id)] = ChatMessageHistory()
    return chat_history[(user_id,conversation_id)]

#创建一个带聊天记录的runable
do_message = RunnableWithMessageHistory(
    chain,
    get_chat_history,
    input_messages_key="input_msg", # 自定义输入的聊天记录的key
    history_messages_key="history_msg", # 对应的MessagesPlaceholder中的variable_name 历史聊天记录
    history_factory_config=[
        ConfigurableFieldSpec(
            id="user_id",
            name="user_id",
            description="用户的唯一标识符",
            default="",
            is_shared=True,
            annotation="",
        ),
        ConfigurableFieldSpec(
            id="conversation_id",
            name="conversation_id",
            description="对话的唯一标识符",
            default="",
            is_shared=True,
            annotation="",
        ),
    ]
)


while True:
    print("-" * 50)
    input_conversation_id = input("请输入对话id:")
    config = {'configurable': {'user_id': 'boss', "conversation_id": input_conversation_id}}
    question = input("请输入你的问题:")
    resp = do_message.invoke({"input_msg":question},config=config)
    print("*" * 10)
    print("回答:",resp.content)
    print("chat_history:",chat_history)
    print("*" * 10)
    # 判断是否继续聊天
    chat_state = input("是否还继续聊天(1-是/0-否)?")
    while chat_state not in ["1","0"]:
        chat_state = input("请只输入Y或者N")
    if chat_state == "0":
        break